package com.csthink.breaker.aspect;

import com.csthink.breaker.annotations.GuavaLimit;
import com.csthink.common.application.ResponseWrapper;
import com.csthink.common.infrastructure.enums.ErrorEnum;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.Objects;

/**
 * 基于 Guava 的 RateLimit 通过 AOP 切面拦截限流注解
 * 仅支持单机版，非分布式限流方案
 *
 * @author <a href="mailto:security.2009@live.cn">Mars</a>
 * @since 2023-09-08
 */
@Slf4j
@Aspect
@Component
public class GuavaLimitAop {

    /**
     * 不同的接口，不同的流量控制
     * key: Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Around("@annotation(com.csthink.breaker.annotations.GuavaLimit)")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        // 获取方法签名上的 GuavaLimit的注解
        GuavaLimit guavaLimit = method.getAnnotation(GuavaLimit.class);

        if (Objects.nonNull(guavaLimit)) {
            // key作用：不同的接口，不同的流量控制
            String key = guavaLimit.key();
            RateLimiter rateLimiter = null;
            // 验证缓存是否有命中key
            if (!limitMap.containsKey(key)) {
                // 创建令牌桶
                rateLimiter = RateLimiter.create(guavaLimit.permitsPerSecond());
                limitMap.put(key, rateLimiter);
                log.info("新建了Guava令牌桶={}，容量={}", key, guavaLimit.permitsPerSecond());
            }

            rateLimiter = limitMap.get(key);
            // 拿令牌
            boolean acquire = rateLimiter.tryAcquire(guavaLimit.timeout(), guavaLimit.timeunit());
            // 拿不到命令，直接返回异常提示
            if (!acquire) {
                log.warn("令牌桶={}，获取令牌失败", key);
                this.responseFail(guavaLimit.msg());
                return null;
            }
        }
        return joinPoint.proceed();
    }

    /**
     * 直接向前端抛出异常
     *
     * @param msg 提示信息
     */
    private void responseFail(String msg) {
        log.warn("responseFail:{}", msg);
        ServletRequestAttributes requestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        if (requestAttributes != null) {
            HttpServletResponse response = requestAttributes.getResponse();
            if (response != null) {
                ResponseWrapper.responseError(ErrorEnum.SERVER_BUSY, msg, response);
            }
        }
    }
}
